/*
 * JGO2RDF creates a RDF representation of GAF2 files, which are automatically sourced from GO ftp site.
 * Copyright (C) 2012 Matthew M. Hindle & Artem Lysenko
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package uk.co.jgo2rdf.goa.parser;

import com.csvreader.CsvReader;
import org.apache.log4j.Logger;
import org.apache.lucene.analysis.KeywordAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.store.NIOFSDirectory;
import org.apache.lucene.util.Version;
import org.xml.sax.SAXException;
import uk.co.jgo2rdf.exception.InvalidValueException;
import uk.co.jgo2rdf.goa.go.OntologyType;
import uk.co.jgo2rdf.goa.model.DefaultGOAEntry;
import uk.co.jgo2rdf.goa.model.GOAEntry;
import uk.co.jgo2rdf.goa.model.GOAFilter;
import uk.co.jgo2rdf.goa.model.MutableGOAEntry;
import uk.co.jgo2rdf.tools.URLIOUtils;

import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathExpressionException;
import javax.xml.xpath.XPathFactory;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.nio.charset.Charset;
import java.text.DateFormat;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.util.*;

import static uk.co.jgo2rdf.goa.go.OntologyType.UnknownOntologyNameException;
import static uk.co.jgo2rdf.goa.go.OntologyType.getOntology;

/**
 * Created by IntelliJ IDEA.
 * User: hindlem
 * Date: 4/8/11
 * Time: 3:18 PM
 */
public class GOAIndexer {

    private static final String[] headers = new String[]{
            "DB", "DBObjectID", "DBObjectSymbol",
            "Qualifier", "GOID",
            "DB:Reference",
            "EvidenceCode",
            "WithFrom",
            "Aspect",
            "DBObjectName",
            "DBObjectSynonym",
            "DBObjectType",
            "Taxon",
            "Date",
            "AssignedBy",
            "AnnotationExtension",
            "GeneProductFormID",
    };
    private DateFormat date;
    private IndexSearcher is;
    private NIOFSDirectory directory;

    /**
     * Represents the Column in the GOA file
     */
    enum Column {
        DB(0), DBObjectID(1), DBObjectSymbol(2),
        Qualifier(3), GOID(4), DBReference(5), EvidenceCode(6),
        WithFrom(7), Aspect(8), DBObjectName(9), DBObjectSynonym(10),
        DBObjectType(11), Taxon(12), Date(13), AssignedBy(14),
        AnnotationExtension(15), GeneProductFormID(16);

        private String name;
        private final int index;

        Column(int columnIndex) {
            this.name = headers[columnIndex];
            this.index = columnIndex;
        }

        public String getName() {
            return name;
        }

        public int getIndex() {
            return index;
        }
    }

    private static final Logger log = Logger.getLogger(GOAIndexer.class);

    private List<InvalidGOAFileException> errors = new ArrayList();

    /**
     * @return errors encountered during indexing
     */
    public List<InvalidGOAFileException> getErrors() {
        return errors;
    }

    public GOAIndexer(String existingIndex, SimpleDateFormat df) throws IOException {
        directory = new NIOFSDirectory(new File(existingIndex));
        is = new IndexSearcher(IndexReader.open(directory, true), null);
        this.date = df;
    }

    /**
     * @param stream2goa the stream to the GOA file
     * @param date       the date format used in the GOA file
     * @param encoding
     * @throws IOException             any problems reading the file
     * @throws InvalidGOAFileException will fail fast if you turned this option on
     */
    public GOAIndexer(InputStream stream2goa,
                      DateFormat date,
                      Charset encoding, String indexDir
    ) throws IOException {
        this.date = date;

        File f = new File(indexDir);
        f.delete();
        f.mkdirs();
        f.mkdir();

        directory = new NIOFSDirectory(f);

        IndexWriterConfig config = new IndexWriterConfig(Version.LUCENE_30,
                new KeywordAnalyzer());
        IndexWriter writer = new IndexWriter(directory, config);
        writer.deleteAll();

        int i = 0; //line counter
        try {

            CsvReader csvReader = new CsvReader(stream2goa,
                    '\t', encoding);
            csvReader.setComment('!');
            csvReader.setUseComments(true);
            csvReader.setTrimWhitespace(true);
            csvReader.setHeaders(headers);

            while (csvReader.readRecord()) {
                i++;
                //System.out.println( csvReader.getColumnCount()+" "+csvReader.getRawRecord());
                if (csvReader.getColumnCount() > 1) { //ignore the empty columns
                    writer.addDocument(createDocument(csvReader));
                }
            }
        } finally {
            writer.prepareCommit();
            writer.commit();
            writer.close(true);
            stream2goa.close();
            log.info("processed " + i + " lines of annotation in file");
        }

        is = new IndexSearcher(IndexReader.open(directory, true), null);
    }

    public Set<GOAEntry> getAllEntries() throws IOException, InvalidGOAFileException, InstantiationException, IllegalAccessException {
        return getAllEntries(null);
    }

    public Set<GOAEntry> getAllEntries(GOAFilter filter) throws IOException, InvalidGOAFileException, InstantiationException, IllegalAccessException {
        final Set<GOAEntry> entries = new HashSet<GOAEntry>();
        IndexReader r = IndexReader.open(directory, true);

        int num = r.numDocs();
        for (int i = 0; i < num; i++) {
            if (!r.isDeleted(i)) {
                Document d = r.document(i);
                GOAEntry entry = createGOAEntry(d);
                if (filter == null || filter.acceptEntry(entry))
                    entries.add(entry);
            }
        }
        r.close();

        return entries;
    }

    public Set<GOAEntry> getEntriesForGene(String gene) throws IOException {
        return getEntriesForGene(gene, null);
    }


    public Set<GOAEntry> getEntriesForSymbol(String gene) throws IOException {
        return getEntriesForSymbol(gene, null);
    }

    /**
     * @return all the entries ever parsed by this Indexer
     */
    public Set<GOAEntry> getEntriesForSymbol(String gene, final GOAFilter filter) throws IOException {

        final Set<GOAEntry> entries = new HashSet<GOAEntry>();

        TermQuery tq = new TermQuery(new Term(Column.DBObjectSymbol.getName(), gene));

        final Collector docCollector = new Collector() {

            private int docBase;

            @Override
            public void setScorer(Scorer scorer) throws IOException {
            }

            @Override
            public void collect(int doc) throws IOException {
                try {
                    GOAEntry entry = createGOAEntry(is.doc(doc + docBase));
                    if (filter == null || filter.acceptEntry(entry))
                        entries.add(entry);
                } catch (Exception e) {
                    log.error(e);
                    e.printStackTrace();
                }
            }

            @Override
            public void setNextReader(IndexReader reader, int docBase) throws IOException {
                this.docBase = docBase;
            }

            @Override
            public boolean acceptsDocsOutOfOrder() {
                return true;
            }
        };

        is.search(tq, docCollector);
        return entries;
    }

    /**
     * @return all the entries ever parsed by this Indexer
     */
    public Set<GOAEntry> getEntriesForGene(String gene, final GOAFilter filter) throws IOException {

        final Set<GOAEntry> entries = new HashSet<GOAEntry>();

        TermQuery tq = new TermQuery(new Term(Column.DBObjectID.getName(), gene));

        final Collector docCollector = new Collector() {

            private int docBase;

            @Override
            public void setScorer(Scorer scorer) throws IOException {
            }

            @Override
            public void collect(int doc) throws IOException {
                try {
                    GOAEntry entry = createGOAEntry(is.doc(doc + docBase));
                    if (filter == null || filter.acceptEntry(entry))
                        entries.add(entry);
                } catch (Exception e) {
                    log.error(e);
                    e.printStackTrace();
                }
            }

            @Override
            public void setNextReader(IndexReader reader, int docBase) throws IOException {
                this.docBase = docBase;
            }

            @Override
            public boolean acceptsDocsOutOfOrder() {
                return true;
            }
        };

        is.search(tq, docCollector);
        return entries;
    }

    private MutableGOAEntry createGOAEntry(Document doc) throws InvalidGOAFileException {
        MutableGOAEntry goaEntryObject = new DefaultGOAEntry();

        try {
            goaEntryObject.setDB(doc.get(Column.DB.getName()).trim());
        } catch (Exception e) {
            log.error(e);
            throw new InvalidGOAFileException(doc.get(Column.DB.getName()).trim(), e, Column.DB);
        }

        try {
            goaEntryObject.setDBObject_ID(doc.get(Column.DBObjectID.getName()).trim());
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.DBObjectID.getName()).trim(), e, Column.DBObjectID);
        }


        try {
            goaEntryObject.setDBObjectSymbol(doc.get(Column.DBObjectSymbol.getName()).trim());
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.DBObjectSymbol.getName()).trim(), e, Column.DBObjectSymbol);
        }

        if (doc.get(Column.Qualifier.getName()) != null) {
            String p_qualifier = doc.get(Column.Qualifier.getName()).trim();
            try {
                if (p_qualifier.length() > 0) //do not assign empty string or anything that's not a go term
                    goaEntryObject.setQualifier(p_qualifier.trim().intern());
            } catch (Exception e) {
                throw new InvalidGOAFileException(doc.get(Column.Qualifier.getName()).trim(), e, Column.Qualifier);
            }
        }

        try {
            String goId = doc.get(Column.GOID.getName()).trim().toUpperCase();
            goaEntryObject.setGO_ID(goId);
        } catch (InvalidValueException e) {
            throw new InvalidGOAFileException(doc.get(Column.GOID.getName()).trim(), e, Column.GOID);
        }


        try {
            goaEntryObject.addDB_References(doc.get(Column.DBReference.getName()).trim().split("\\|"));
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.DBReference.getName()).trim(), e, Column.DBReference);
        }

        try {
            goaEntryObject.setEvidenceCode(doc.get(Column.EvidenceCode.getName()).trim());
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.EvidenceCode.getName()).trim(), e, Column.EvidenceCode);
        }

        try {
            goaEntryObject.setWithFrom(doc.get(Column.WithFrom.getName()).trim());
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.WithFrom.getName()).trim(), e, Column.WithFrom);
        }

        String ontology = doc.get(Column.Aspect.getName()).trim();
        try {
            goaEntryObject.setAspect(getOntology(ontology));
        } catch (UnknownOntologyNameException e) {
            try {
                goaEntryObject.setAspect(getAspectForTerm(goaEntryObject.getGO_ID_String()));
            } catch (Exception e1) {
                throw new InvalidGOAFileException(ontology, e, Column.Aspect);
            }
        } catch (Exception e) {
            throw new InvalidGOAFileException(ontology, e, Column.Aspect);
        }

        try {
            goaEntryObject.setDBObjectName(doc.get(Column.DBObjectName.getName()).trim());
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.DBObjectName.getName()).trim(), e, Column.DBObjectName);
        }


        String pdbObjectSynonym = doc.get(Column.DBObjectSynonym.getName()).trim();
        if (pdbObjectSynonym.length() > 0)  //do not assign empty string
            try {
                goaEntryObject.setDBObjectSynonym(pdbObjectSynonym);
            } catch (Exception e) {
                throw new InvalidGOAFileException(doc.get(Column.DBObjectSynonym.getName()).trim(), e, Column.DBObjectSynonym);
            }

        try {
            goaEntryObject.setDBObjectType(doc.get(Column.DBObjectType.getName()).trim());
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.DBObjectType.getName()).trim(), e, Column.DBObjectType);
        }

        String[] taxonsString = doc.get(Column.Taxon.getName()).toLowerCase().replaceAll("taxon:", "").trim().split("\\|");
        int[] taxons = new int[taxonsString.length];
        try {
            for (int i = 0; i < taxonsString.length; i++)
                taxons[i] = Integer.parseInt(taxonsString[i]);
        } catch (NumberFormatException e) {
            throw new InvalidGOAFileException(Arrays.toString(taxonsString), e, Column.Taxon);
        }
        if (taxons.length == 0)
            throw new InvalidGOAFileException(Arrays.toString(taxonsString), "no species taxonomy defined", Column.Taxon);

        Arrays.sort(taxons); //this makes equivilence comparison easier for finding duplicate entries

        try {
            goaEntryObject.setTaxons(taxons);
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.Taxon.getName()).trim(), e, Column.Taxon);
        }

        if (doc.get(Column.Date.getName()) != null)
            try {
                goaEntryObject.setDate(date.parse(doc.get(Column.Date.getName()).trim()));
            } catch (ParseException e) {
                throw new InvalidGOAFileException(doc.get(Column.Date.getName()).trim(), e, Column.Date);
            } catch (Exception e) {
                throw new InvalidGOAFileException(doc.get(Column.Date.getName()).trim(), e, Column.Date);
            }

        try {
            goaEntryObject.setAssignedBy(doc.get(Column.AssignedBy.getName()).trim());
        } catch (Exception e) {
            throw new InvalidGOAFileException(doc.get(Column.AssignedBy.getName()).trim(), e, Column.AssignedBy);
        }

        if (doc.get(Column.AnnotationExtension.getName()) != null) {
            String p_annotationExtension = doc.get(Column.AnnotationExtension.getName()).trim();
            if (p_annotationExtension.length() > 0) //do not assign empty string
                try {
                    goaEntryObject.setAnnotation_Extension(p_annotationExtension);
                } catch (Exception e) {
                    throw new InvalidGOAFileException(doc.get(Column.AnnotationExtension.getName()).trim(), e, Column.AnnotationExtension);
                }

        }

        if (doc.get(Column.GeneProductFormID.getName()) != null) {
            String p_geneproductFormId = doc.get(Column.GeneProductFormID.getName()).trim();
            if (p_geneproductFormId.length() > 0) //do not assign empty string
                try {
                    goaEntryObject.setGeneProductFormID(p_geneproductFormId);
                } catch (Exception e) {
                    throw new InvalidGOAFileException(doc.get(Column.GeneProductFormID.getName()).trim(), e, Column.GeneProductFormID);
                }
        }
        return goaEntryObject;
    }

    private Document createDocument(CsvReader reader) throws IOException {

        Document doc = new Document();
        doc.add(new Field(Column.DB.getName(),
                reader.get(Column.DB.getName()).trim(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));
        doc.add(new Field(Column.DBObjectID.getName(),
                reader.get(Column.DBObjectID.getName()).trim(),
                Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));
        doc.add(new Field(Column.DBObjectSymbol.getName(),
                reader.get(Column.DBObjectSymbol.getName()).trim(),
                Field.Store.YES, Field.Index.ANALYZED_NO_NORMS));

        String p_qualifier = reader.get(Column.Qualifier.getName()).trim();
        if (p_qualifier.length() > 0) //do not assign empty string or anything that's not a go term
            doc.add(new Field(Column.Qualifier.getName(),
                    reader.get(Column.Qualifier.getName()).trim(),
                    Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.GOID.getName(),
                reader.get(Column.GOID.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.DBReference.getName(),
                reader.get(Column.DBReference.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.EvidenceCode.getName(),
                reader.get(Column.EvidenceCode.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.WithFrom.getName(),
                reader.get(Column.WithFrom.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.Aspect.getName(),
                reader.get(Column.Aspect.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.DBObjectName.getName(),
                reader.get(Column.DBObjectName.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.DBObjectSynonym.getName(),
                reader.get(Column.DBObjectSynonym.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.DBObjectType.getName(),
                reader.get(Column.DBObjectType.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.Taxon.getName(),
                reader.get(Column.Taxon.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.Date.getName(),
                reader.get(Column.Date.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.AssignedBy.getName(),
                reader.get(Column.AssignedBy.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.AnnotationExtension.getName(),
                reader.get(Column.AnnotationExtension.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        doc.add(new Field(Column.GeneProductFormID.getName(),
                reader.get(Column.GeneProductFormID.getName()).trim().toUpperCase(),
                Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));

        return doc;
    }

    public final Map<String, OntologyType> cache = new HashMap<String, OntologyType>();

    public OntologyType getAspectForTerm(String term) throws IOException, ParserConfigurationException, SAXException, XPathExpressionException, UnknownOntologyNameException, InterruptedException {
        if (cache.containsKey(term.toUpperCase())) {
            return cache.get(term.toUpperCase());
        }
        log.info("asking EBI webservice what aspect " + term + " belongs to");

        // URL a GO Term in OBO xml format
        URL url = new URL("http://www.ebi.ac.uk/QuickGO/GTerm?id=" + term.toUpperCase() + "&format=oboxml");
        // Connect
        byte[] bytes = URLIOUtils.getByteArray(url, 10);

        org.w3c.dom.Document xml = DocumentBuilderFactory.newInstance().newDocumentBuilder().parse(new ByteArrayInputStream(bytes));

        // XPath is here used to locate parts of an XML document
        XPath xpath = XPathFactory.newInstance().newXPath();
        OntologyType ontology = getOntology(xpath.compile("/obo/term/namespace").evaluate(xml));
        cache.put(term.toUpperCase(), ontology);
        return ontology;
    }


    public static class InvalidGOAFileException extends Exception {

        InvalidGOAFileException(Exception e) {
            super("GOA file is invalid", e);
        }

        InvalidGOAFileException(String item, Exception e, Column column) {
            super("GOA file is invalid at column " + column.getIndex() + " which should be the field " + column.getName() + ", entry reads \"" + item + "\"", e);
        }

        InvalidGOAFileException(String item, String message, Column column) {
            super("GOA file is invalid at column " + column.getIndex() + " which should be the field " + column.getName() + ", entry reads \"" + item + "\"" + ": " + message);
        }

    }

}
